import os
import sys
import struct
from elftools.elf.elffile import ELFFile

def read_map_file(map_file_path):
    if not os.path.exists(map_file_path):
        print(f"Map file {map_file_path} does not exist.")
        sys.exit(1)
    mapping = {}
    with open(map_file_path, 'r') as f:
        for line in f.readlines():
            line = line.strip()
            if line:
                try:
                    a, b = line.split(",")
                    if len(a) < len(b):
                        print(f"In map file, A string length must be greater than or equal to B. Got {a} and {b}.")
                        sys.exit(1)
                    mapping[a] = b
                except ValueError:
                    print(f"Map file format error at line: {line}.")
                    sys.exit(1)
    return mapping

def process_elf_file(elf_file_path, mapping):
    if not os.path.exists(elf_file_path):
        print(f"ELF file {elf_file_path} does not exist.")
        sys.exit(1)
    try:
        with open(elf_file_path, 'rb+') as f:
            elffile = ELFFile(f)
            rodata_section = None
            for section in elffile.iter_sections():
                if section.name == '.rodata':
                    rodata_section = section
                    break
            if rodata_section is None:
                print("No .rodata section found.")
                sys.exit(1)
            rodata = rodata_section.data()
            for a, b in mapping.items():
                rodata = rodata.replace(a.encode(), b.encode().ljust(len(a), b'\0'))
            f.seek(rodata_section['sh_offset'])
            f.write(rodata)
    except Exception as e:
        print(f"Error while processing ELF file: {e}.")
        sys.exit(1)

if __name__ == "__main__":
    if len(sys.argv) < 3:
        print("Please provide the ELF file path and map file path as arguments.")
        sys.exit(1)
    elf_file_path = sys.argv[1]
    map_file_path = sys.argv[2]
    mapping = read_map_file(map_file_path)
    process_elf_file(elf_file_path, mapping)
